package main

import (
	"crypto/ecdsa"
	"device/util"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/gogf/gf/v2/net/gtcp"
	"time"

	"device/model"
)

type Client struct {
	ipaddr string
	token  string
}

type GataWay struct {
	ipaddr string
}

func main() {
	s := flag.String("ip", "127.0.0.1:8990", "密钥中心ip")
	etcdif := flag.Bool("e", false, "是否开启etcd")
	flag.Parse()
	fmt.Println("KeyCenter Listen AT", *s)

	NewServer(*s, *etcdif).Run()
}

func NewServer(ip string, etcdif bool) *gtcp.Server {
	if etcdif {
		registerKeyCenterServer(ip)
	}
	//缓存client未来拿的证书
	clientWaitingMap := make(map[string]*ecdsa.PrivateKey)
	//GataWayWaitingMap := make(map[string]*ecdsa.PrivateKey)
	server := gtcp.NewServer(ip, func(conn *gtcp.Conn) {
		conn.SetDeadline(time.Now().Add(10 * time.Second))
		defer conn.Close()
		for {
			//log.Println("新连接", conn.RemoteAddr())
			pkg, err := conn.Recv(-1)
			//log.Println("收到包！！", string(pkg), "包结束")
			if err != nil {
				fmt.Println(err)
			}
			msg := new(model.Message)
			err = json.Unmarshal(pkg, msg)

			if err != nil {
				conn.Send([]byte("json err!"))
			}
			switch msg.Type {
			case model.KeyGenerate: //生成密钥请求，直接返回，原则上不应该有错误处理
				device := msg.Msg.(map[string]interface{})
				fmt.Println("断言成功?")
				Key := GenerateKey()
				fmt.Println("生成证书完成")
				clientWaitingMap[device["Token"].(string)] = Key
				//go deleteMapKey(device["Token"].(string), clientWaitingMap) //5s客户端没来要证书就删除证书缓存
				_, bytes := sm2ToBytes(Key) //直接是证书的格式
				fmt.Println("返回给网关的证书", string(bytes))
				conn.Send(bytes)
				break
			case model.ClientPem: //客户端来要密钥
				fmt.Println(msg)
				message := msg.Msg.(map[string]interface{})
				for i := 0; i < 5; i++ { //客户端来要密钥，循环等待5s,每秒检测一次密钥是否被生成，防止因为密钥生成的慢而被误判
					key, ok := clientWaitingMap[message["Token"].(string)]
					if ok {
						pribytes, pubytes := sm2ToBytes(key)
						clientKey := new(model.Key)
						clientKey.PRIKEY = model.PrivateKey{Data: pribytes}
						clientKey.PUBKEY = model.PublicKey{Data: pubytes}
						marshal, _ := json.Marshal(clientKey)
						conn.Send(marshal)
						fmt.Println("返回给客户端的证书：", string(marshal))
						conn.Close()
						break
						delete(clientWaitingMap, message["Token"].(string))
					} else {
						fmt.Println("键不存在")
					}
					time.Sleep(2 * time.Second)
				}
				conn.SendPkg([]byte("超时未连接"))
				break
			}
			break
		}
	})
	return server
}

func registerKeyCenterServer(ipaddr string) {
	util.New("/etcd/keycenter", ipaddr)
}

func deleteMapKey(token string, clientWaitingMap map[string]*ecdsa.PrivateKey) {
	time.Sleep(10 * time.Second)
	fmt.Println("删除证书")
	_, ok := clientWaitingMap[token]
	if ok {
		delete(clientWaitingMap, token)
	}
}
